package com.learning.datastruct.tree;

import java.util.Comparator;

/**
 * @author ADMIN
 * @Classname AVLTree
 * @description
 * @date 2022/1/24 5:14
 */
@SuppressWarnings("AlibabaClassNamingShouldBeCamel")
public class AVLTree<E> extends BalancedBinarySearchTree<E> {
    public AVLTree() {
        super();
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    private static class AVLNode<E> extends Node<E> {
        protected AVLNode<E> left;
        protected AVLNode<E> right;
        protected AVLNode<E> parent;
        private int height = 1;
        protected AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }
        private int balanceFactor() {
            int leftHeight = this.left == null ? 0 : this.left.height;
            int rightHeight = this.right == null ? 0 : this.right.height;
            return leftHeight - rightHeight;
        }
        private void updateHeight() {
            int leftHeight = this.left == null ? 0 : this.left.height;
            int rightHeight = this.right == null ? 0 : this.right.height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }
        private boolean isBalanced() {
            return Math.abs(this.balanceFactor()) <= 1;
        }
        private AVLNode<E> tallerChild() {
            int leftHeight = this.left == null ? 0 : this.left.height;
            int rightHeight = this.right == null ? 0 : this.right.height;
            if (leftHeight > rightHeight) {
                return left;
            } else if (leftHeight < rightHeight) {
                return right;
            } else {
                return isLeftChild() ? left : right;
            }
        }
    }

    @Override
    protected void afterAdd(Node<E> newNode) {
        while ((newNode = newNode.parent) != null) {
            AVLNode<E> node = (AVLNode<E>) newNode;
            if (node.isBalanced()) {
                updateHeight(node);
            } else {
                reBalance(node);
                break;
            }
        }
    }

    @Override
    protected void afterRemove(BinaryTree.Node<E> newNode) {
        while ((newNode = newNode.parent) != null) {
            AVLNode<E> node = (AVLNode<E>) newNode;
            if (node.isBalanced()) {
                updateHeight(node);
            } else {
                reBalance(node);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    private void reBalance(AVLNode<E> grand) {
        AVLNode<E> parent = grand.tallerChild();
        AVLNode<E> node = parent.tallerChild();
        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                // LL
                rotateRight(grand);
            } else {
                // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        } else {
            if (node.isLeftChild()) {
                // RL
                rotateRight(parent);
                rotateLeft(grand);
            } else {
                // RR
                rotateLeft(grand);
            }
        }
    }

    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> child) {
        super.afterRotate(grand, parent, child);
        updateHeight(grand);
        updateHeight(parent);
    }

    private void updateHeight(Node<E> node){
        ((AVLNode<E>)node).updateHeight();
    }

}
